/****************************************************************************
 *
 * Copyright 2016 Samsung Electronics All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the specific
 * language governing permissions and limitations under the License.
 *
 ****************************************************************************/
/*
 * Wslay - The WebSocket Library
 *
 * Copyright (c) 2011, 2012 Tatsuhiro Tsujikawa
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
 * LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 * OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */
#include "wslay_frame.h"

#include <stddef.h>
#include <string.h>
#include <assert.h>

#include "wslay_net.h"

#define wslay_min(A, B) (((A) < (B)) ? (A) : (B))

int wslay_frame_context_init(wslay_frame_context_ptr *ctx, const struct wslay_frame_callbacks *callbacks, void *user_data)
{
	*ctx = (wslay_frame_context_ptr)malloc(sizeof(struct wslay_frame_context));
	if (*ctx == NULL) {
		return -1;
	}
	memset(*ctx, 0, sizeof(struct wslay_frame_context));
	(*ctx)->istate = RECV_HEADER1;
	(*ctx)->ireqread = 2;
	(*ctx)->ostate = PREP_HEADER;
	(*ctx)->user_data = user_data;
	(*ctx)->ibufmark = (*ctx)->ibuflimit = (*ctx)->ibuf;
	(*ctx)->callbacks = *callbacks;
	return 0;
}

void wslay_frame_context_free(wslay_frame_context_ptr ctx)
{
	free(ctx);
}

ssize_t wslay_frame_send(wslay_frame_context_ptr ctx, struct wslay_frame_iocb *iocb)
{
	if (iocb->data_length > iocb->payload_length) {
		return WSLAY_ERR_INVALID_ARGUMENT;
	}
	if (ctx->ostate == PREP_HEADER) {
		uint8_t *hdptr = ctx->oheader;
		memset(ctx->oheader, 0, sizeof(ctx->oheader));
		*hdptr |= (iocb->fin << 7) & 0x80u;
		*hdptr |= (iocb->rsv << 4) & 0x70u;
		*hdptr |= iocb->opcode & 0xfu;
		++hdptr;
		*hdptr |= (iocb->mask << 7) & 0x80u;
		if (wslay_is_ctrl_frame(iocb->opcode) && iocb->payload_length > 125) {
			return WSLAY_ERR_INVALID_ARGUMENT;
		}
		if (iocb->payload_length < 126) {
			*hdptr |= iocb->payload_length;
			++hdptr;
		} else if (iocb->payload_length < (1 << 16)) {
			uint16_t len = htons(iocb->payload_length);
			*hdptr |= 126;
			++hdptr;
			memcpy(hdptr, &len, 2);
			hdptr += 2;
		} else if (iocb->payload_length < (1ull << 63)) {
			uint64_t len = hton64(iocb->payload_length);
			*hdptr |= 127;
			++hdptr;
			memcpy(hdptr, &len, 8);
			hdptr += 8;
		} else {
			/* Too large payload length */
			return WSLAY_ERR_INVALID_ARGUMENT;
		}
		if (iocb->mask) {
			if (ctx->callbacks.genmask_callback(ctx->omaskkey, 4, ctx->user_data) != 0) {
				return WSLAY_ERR_INVALID_CALLBACK;
			} else {
				ctx->omask = 1;
				memcpy(hdptr, ctx->omaskkey, 4);
				hdptr += 4;
			}
		}
		ctx->ostate = SEND_HEADER;
		ctx->oheadermark = ctx->oheader;
		ctx->oheaderlimit = hdptr;
		ctx->opayloadlen = iocb->payload_length;
		ctx->opayloadoff = 0;
	}
	if (ctx->ostate == SEND_HEADER) {
		ptrdiff_t len = ctx->oheaderlimit - ctx->oheadermark;
		ssize_t r;
		int flags = 0;
		if (iocb->data_length > 0) {
			flags |= WSLAY_MSG_MORE;
		};
		r = ctx->callbacks.send_callback(ctx->oheadermark, len, flags, ctx->user_data);
		if (r > 0) {
			if (r > len) {
				return WSLAY_ERR_INVALID_CALLBACK;
			} else {
				ctx->oheadermark += r;
				if (ctx->oheadermark == ctx->oheaderlimit) {
					ctx->ostate = SEND_PAYLOAD;
				} else {
					return WSLAY_ERR_WANT_WRITE;
				}
			}
		} else {
			return WSLAY_ERR_WANT_WRITE;
		}
	}
	if (ctx->ostate == SEND_PAYLOAD) {
		size_t totallen = 0;
		if (iocb->data_length > 0) {
			if (ctx->omask) {
				uint8_t temp[4096];
				const uint8_t *datamark = iocb->data, *datalimit = iocb->data + iocb->data_length;
				while (datamark < datalimit) {
					size_t datalen = datalimit - datamark;
					const uint8_t *writelimit = datamark + wslay_min(sizeof(temp), datalen);
					size_t writelen = writelimit - datamark;
					ssize_t r;
					size_t i;
					for (i = 0; i < writelen; ++i) {
						temp[i] = datamark[i] ^ ctx->omaskkey[(ctx->opayloadoff + i) % 4];
					}
					r = ctx->callbacks.send_callback(temp, writelen, 0, ctx->user_data);
					if (r > 0) {
						if ((size_t)r > writelen) {
							return WSLAY_ERR_INVALID_CALLBACK;
						} else {
							datamark += r;
							ctx->opayloadoff += r;
							totallen += r;
						}
					} else {
						if (totallen > 0) {
							break;
						} else {
							return WSLAY_ERR_WANT_WRITE;
						}
					}
				}
			} else {
				ssize_t r;
				r = ctx->callbacks.send_callback(iocb->data, iocb->data_length, 0, ctx->user_data);
				if (r > 0) {
					if ((size_t)r > iocb->data_length) {
						return WSLAY_ERR_INVALID_CALLBACK;
					} else {
						ctx->opayloadoff += r;
						totallen = r;
					}
				} else {
					return WSLAY_ERR_WANT_WRITE;
				}
			}
		}
		if (ctx->opayloadoff == ctx->opayloadlen) {
			ctx->ostate = PREP_HEADER;
		}
		return totallen;
	}
	return WSLAY_ERR_INVALID_ARGUMENT;
}

static void wslay_shift_ibuf(wslay_frame_context_ptr ctx)
{
	ptrdiff_t len = ctx->ibuflimit - ctx->ibufmark;
	memmove(ctx->ibuf, ctx->ibufmark, len);
	ctx->ibuflimit = ctx->ibuf + len;
	ctx->ibufmark = ctx->ibuf;
}

static ssize_t wslay_recv(wslay_frame_context_ptr ctx)
{
	ssize_t r;
	if (ctx->ibufmark != ctx->ibuf) {
		wslay_shift_ibuf(ctx);
	}
	r = ctx->callbacks.recv_callback(ctx->ibuflimit, ctx->ibuf + sizeof(ctx->ibuf) - ctx->ibuflimit, 0, ctx->user_data);
	if (r > 0) {
		ctx->ibuflimit += r;
	} else {
		r = WSLAY_ERR_WANT_READ;
	}
	return r;
}

#define WSLAY_AVAIL_IBUF(ctx) ((size_t)(ctx->ibuflimit - ctx->ibufmark))

ssize_t wslay_frame_recv(wslay_frame_context_ptr ctx, struct wslay_frame_iocb *iocb)
{
	ssize_t r;
	if (ctx->istate == RECV_HEADER1) {
		uint8_t fin, opcode, rsv, payloadlen;
		if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) {
			if ((r = wslay_recv(ctx)) <= 0) {
				return r;
			}
		}
		if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) {
			return WSLAY_ERR_WANT_READ;
		}
		fin = (ctx->ibufmark[0] >> 7) & 1;
		rsv = (ctx->ibufmark[0] >> 4) & 7;
		opcode = ctx->ibufmark[0] & 0xfu;
		ctx->iom.opcode = opcode;
		ctx->iom.fin = fin;
		ctx->iom.rsv = rsv;
		++ctx->ibufmark;
		ctx->imask = (ctx->ibufmark[0] >> 7) & 1;
		payloadlen = ctx->ibufmark[0] & 0x7fu;
		++ctx->ibufmark;
		if (wslay_is_ctrl_frame(opcode) && (payloadlen > 125 || !fin)) {
			return WSLAY_ERR_PROTO;
		}
		if (payloadlen == 126) {
			ctx->istate = RECV_EXT_PAYLOADLEN;
			ctx->ireqread = 2;
		} else if (payloadlen == 127) {
			ctx->istate = RECV_EXT_PAYLOADLEN;
			ctx->ireqread = 8;
		} else {
			ctx->ipayloadlen = payloadlen;
			ctx->ipayloadoff = 0;
			if (ctx->imask) {
				ctx->istate = RECV_MASKKEY;
				ctx->ireqread = 4;
			} else {
				ctx->istate = RECV_PAYLOAD;
			}
		}
	}
	if (ctx->istate == RECV_EXT_PAYLOADLEN) {
		if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) {
			if ((r = wslay_recv(ctx)) <= 0) {
				return r;
			}
			if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) {
				return WSLAY_ERR_WANT_READ;
			}
		}
		ctx->ipayloadlen = 0;
		ctx->ipayloadoff = 0;
		memcpy((uint8_t *)&ctx->ipayloadlen + (8 - ctx->ireqread), ctx->ibufmark, ctx->ireqread);
		ctx->ipayloadlen = ntoh64(ctx->ipayloadlen);
		ctx->ibufmark += ctx->ireqread;
		if (ctx->ireqread == 8) {
			if (ctx->ipayloadlen < (1 << 16) || ctx->ipayloadlen & (1ull << 63)) {
				return WSLAY_ERR_PROTO;
			}
		} else if (ctx->ipayloadlen < 126) {
			return WSLAY_ERR_PROTO;
		}
		if (ctx->imask) {
			ctx->istate = RECV_MASKKEY;
			ctx->ireqread = 4;
		} else {
			ctx->istate = RECV_PAYLOAD;
		}
	}
	if (ctx->istate == RECV_MASKKEY) {
		if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) {
			if ((r = wslay_recv(ctx)) <= 0) {
				return r;
			}
			if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) {
				return WSLAY_ERR_WANT_READ;
			}
		}
		memcpy(ctx->imaskkey, ctx->ibufmark, 4);
		ctx->ibufmark += 4;
		ctx->istate = RECV_PAYLOAD;
	}
	if (ctx->istate == RECV_PAYLOAD) {
		uint8_t *readlimit, *readmark;
		uint64_t rempayloadlen = ctx->ipayloadlen - ctx->ipayloadoff;
		if (WSLAY_AVAIL_IBUF(ctx) == 0 && rempayloadlen > 0) {
			if ((r = wslay_recv(ctx)) <= 0) {
				return r;
			}
		}
		readmark = ctx->ibufmark;
		readlimit = WSLAY_AVAIL_IBUF(ctx) < rempayloadlen ? ctx->ibuflimit : ctx->ibufmark + rempayloadlen;
		if (ctx->imask) {
			for (; ctx->ibufmark != readlimit; ++ctx->ibufmark, ++ctx->ipayloadoff) {
				ctx->ibufmark[0] ^= ctx->imaskkey[ctx->ipayloadoff % 4];
			}
		} else {
			ctx->ibufmark = readlimit;
			ctx->ipayloadoff += readlimit - readmark;
		}
		iocb->fin = ctx->iom.fin;
		iocb->rsv = ctx->iom.rsv;
		iocb->opcode = ctx->iom.opcode;
		iocb->payload_length = ctx->ipayloadlen;
		iocb->mask = ctx->imask;
		iocb->data = readmark;
		iocb->data_length = ctx->ibufmark - readmark;
		if (ctx->ipayloadlen == ctx->ipayloadoff) {
			ctx->istate = RECV_HEADER1;
			ctx->ireqread = 2;
		}
		return iocb->data_length;
	}
	return WSLAY_ERR_INVALID_ARGUMENT;
}
